"""
基于开盘收盘差异的多标的策略
使用分钟K线，多品种RSRS策略
"""

import numpy as np
import pandas as pd
from datetime import datetime
from tqsdk import TqApi, TargetPosTask
from sklearn.linear_model import LinearRegression
from .base_strategy import BaseStrategy

class OcDiffStrategy(BaseStrategy):
    strategy_name = "oc_diff"
    
    """
    基于开盘收盘差异指标的多标的策略
    使用分钟K线，多品种策略
    """
    def __init__(self, api: TqApi, symbol_info:dict, market_period:int):
        super().__init__(api, symbol_info, market_period)
        for symbol,info in symbol_info.items():
            self.all_info[symbol]["corr_series"] = []

        
    def calculate_signal(self, symbol: str,current_pos:int, kline: pd.DataFrame) -> float:
        """
        计算交易信号
        
        Args:
            symbol: 品种代码
            current_pos: 当前持仓
            kline: K线数据
            
        Returns:
            交易信号: 1-做多, -1-做空, 0-无信号
        """
        score = self.calculate_oc_corr(symbol, kline)
        score_period = self.all_info[symbol]["score_period"]

        close_mean = np.mean(kline.close.values[-score_period:])
        open_mean = np.mean(kline.open.values[-score_period:])
        
        long_threshold = self.all_info[symbol]["long_threshold"]
        short_threshold = self.all_info[symbol]["short_threshold"]
        close_threshold = self.all_info[symbol]["close_threshold"]
        if current_pos == 0:
            if score > long_threshold and close_mean > open_mean:
                signal =  1  # 做多信号
            elif score < short_threshold and close_mean < open_mean:
                signal = -1  # 做空信号
            else:
                signal = 0  # 无信号
        elif current_pos > 0:
            if score < short_threshold and close_mean < open_mean:
                signal = -1
            elif score < close_threshold:
                signal = 0  # 平多信号
            else:
                signal = 1  # 无信号
        elif current_pos < 0:
            if score > long_threshold and close_mean > open_mean:
                signal = 1  # 做多信号
            elif score > -close_threshold:
                signal = 0  # 平空信号
            else:
                signal = -1
        return signal, score

    
    def calculate_oc_corr(self, symbol, kline) -> float:
        """
        计算RSRS指标
        
        Args:
            symbol: 品种代码
            kline: K线数据
            
        Returns:
            标准化后的RSRS值
        """
        score_period = self.all_info[symbol]["score_period"]
        std_period = self.all_info[symbol]["std_period"]
        
        # 获取最近的数据
        close_prices = kline.close.values
        open_prices = kline.open.values
        
        corr_series = self.all_info[symbol]["corr_series"]

        # 首次计算，从足够的历史数据开始计算
        if len(corr_series) < std_period:
            for i in range(score_period, len(close_prices)+1):
                # 获取计算窗口内的高低价
                window_close = close_prices[i-score_period:i]
                window_open = open_prices[i-score_period:i]
                mem = (window_close - window_open).mean()
                corr = np.corrcoef(window_close, window_open)[0][1]*(1 if mem>=0 else -1)
                self.all_info[symbol]["corr_series"].append(corr)
        # 后续更新计算
        else:
            window_close = close_prices[-score_period:]
            window_open = open_prices[-score_period:]
            mem = (window_close - window_open).mean()
            corr = np.corrcoef(window_close, window_open)[0][1]*(1 if mem>=0 else -1)
            self.all_info[symbol]["corr_series"][:-1] = self.all_info[symbol]["corr_series"][1:]
            self.all_info[symbol]["corr_series"][-1]=corr
        
        # 获取最新的RSRS斜率
        current_score = self.all_info[symbol]["corr_series"][-1]
        
        # 标准化处理
        recent_score = self.all_info[symbol]["corr_series"][-std_period:] 
        mean_score = np.mean(recent_score)
        std_score = np.std(recent_score)
        
        if std_score > 0:
            score = (current_score - mean_score) / std_score
        else:
            score = 0.0
        
        return score
